## Helper functions for the processing steps of
##
## Christopher Adolph, Kenya Amano, Bree Bang-Jensen, 
## Nancy Fullman, Beatrice Magistro, Grace Reinke, Rachel Castellano,## Megan Erickson, and John Wilkerson, "The Pandemic Policy U-Turn:## The role of partisanship, public health, and race in decisions
## to ease COVID-19 social distancing policies in the U.S.", 
## forthcoming in Perspectives on Politics.
## 
## Kenya Amano
## 18 June 2021



state.abb.add <-
  c(state.abb,
    "DC",  # District of Columbia	
    "PR"  # Puerto Rico
  )

state.name.add <-
  c(state.name,
    "District of Columbia",
    "Puerto Rico"   
  )


# convert state abb to state name    
abb2name <- 
  function(x){
    
    state.name.add[grep(x, state.abb.add)] %>% 
      as.character()
  }

# convert state name to state abb    
name2abb <- 
  function(x){
    if(x == "Virginia"){"VA"}
    else{state.abb.add[grep(x, state.name.add)] %>% 
        as.character()}
  }

### Age bracket var name function


age_bracket_name_fun <- 
  function(x){

to <- "to"

from_age <- 
  str_extract(x,"[:digit:]+(?=\\b to \\b)")
to_age <- 
  str_extract(x,"(?<=\\b to \\b)[:digit:]+")

paste0("Pop", from_age, "to", to_age)
  }


age_bracket_fun <-
  function(x){
    if(x == "All Ages"){"Pop"}
    else if(x == "Under 5") {"PopUnder5"}
    else if(x == "95 plus"){"PopOver95"}
    else if(x == "1 to 4"){"delete"}
    else if(x == "Post Neonatal"){"delete"}
    else if(x == "Late Neonatal"){"delete"}
    else if(x == "Early Neonatal"){"delete"}
    else if(x == "Birth"){"delete"}
    else {age_bracket_name_fun(x)}
  }



###################################
lag1 <- function(x){
  lag(x, 1)}

lag2 <- function(x){
  lag(x, 2)}

lag3 <- function(x){
  lag(x, 3)}

lag4 <- function(x){
  lag(x, 4)}

lag5 <- function(x){
  lag(x, 5)}

lag6 <- function(x){
  lag(x, 6)}

lag7 <- function(x){
  lag(x, 7)}

lag8 <- function(x){
  lag(x, 8)}

lag9 <- function(x){
  lag(x, 9)}

lag10 <- function(x){
  lag(x, 10)}

lag11 <- function(x){
  lag(x, 11)}

lag12 <- function(x){
  lag(x, 12)}

lag13 <- function(x){
  lag(x, 13)}

lag14 <- function(x){
  lag(x, 14)}



Create2DaysMAVG <- 
  function(x){
    (x+lag(x,1))/2}

Create4DaysMAVG <- 
  function(x){
    (x+lag(x,1)+lag(x,2)+lag(x,3))/4}


Create7DaysMAVG <- 
  function(x){
    (x+lag(x,1)+lag(x,2)+lag(x,3)+lag(x,4)+lag(x,5)+lag(x,6))/7}

Create14DaysMAVG <- 
  function(x){
    (x+lag(x,1)+lag(x,2)+lag(x,3)+lag(x,4)+lag(x,5)+lag(x,6)+
     lag(x,7)+lag(x,8)+lag(x,9)+lag(x,10)+lag(x,11)+lag(x,12)+lag(x,13))/14}



##################################################
CreateLagDF <- 
  function(df, vars){
    nlags <- 14:1
    df_lags <- map(nlags, ~lag(df %>% select(vars), .x)) %>%
      as.data.frame() %>%
      setNames(paste0("lag_", nlags))
    
    Newdf <- bind_cols(df %>% select(state_name, date_num, vars), df_lags)
    
  }

CreateSlopeLM <- 
  function(Lag14, Lag13, Lag12, Lag11, Lag10, Lag9, Lag8, Lag7, Lag6, Lag5, Lag4, Lag3, Lag2, Lag1, RegModel){
    dta <- 
      tibble(Ydta = c(Lag14, Lag13, Lag12, Lag11, Lag10, Lag9, Lag8, Lag7, Lag6, Lag5, Lag4, Lag3, Lag2, Lag1),
             Xdta = c(1:14))
    
    NACheck <- 
      dta %>% filter(is.na(Ydta)==T) %>% nrow()
    
    if(NACheck ==0){
      if(RegModel == "LM"){
        res <- lm(Ydta ~ Xdta, data= dta)
        coef <- coef(res)[2] %>% as.numeric() 
      }
      if(RegModel == "RobustLM"){
        if(mean(dta$Ydta == 0)){
          coef <- 0
        }else{
          res <- MASS::rlm(Ydta ~ Xdta, data= dta, method = "MM")
          coef <- coef(res)[2] %>% as.numeric()  
        }
      }
    }else(coef <- NA)
    
    coef
  }


CreateSlopeDF <- 
  function(DF, vars, Model){
    
      pmap_dfr(list(df = DF %>% ungroup() %>% split(.$state_name),
                    vars = vars),
               CreateLagDF) %>% 
      
      mutate(LagSlope = pmap(list(Lag14= lag_14,
                                  Lag13= lag_12,
                                  Lag12= lag_12,
                                  Lag11= lag_11,
                                  Lag10= lag_10,
                                  Lag9= lag_9,
                                  Lag8= lag_8,
                                  Lag7= lag_7,
                                  Lag6= lag_6,
                                  Lag5= lag_5,
                                  Lag4= lag_4,
                                  Lag3= lag_3,
                                  Lag2= lag_2,
                                  Lag1= lag_1,
                                  RegModel = Model),
                             CreateSlopeLM)) %>% 
      unnest(LagSlope) %>% 
      select(state_name, date_num, LagSlope) 
    
  }

##########################
# tidyverse filter helper


`%!in%` <- Negate(`%in%`) 

